import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.preprocessing import MinMaxScaler
import tensorflow as tf
from tensorflow.keras.models import load_model
import talib
import warnings

# 忽略警告
warnings.filterwarnings('ignore')

def add_technical_indicators(df):
    """添加技术指标作为特征，处理异常情况"""
    # 确保数据量足够
    if len(df) < 50:
        print(f"警告: 数据量较少 ({len(df)} 行), 某些指标可能无法计算")

    # 确保必要的价格和成交量列是数值类型
    numeric_cols = ['open', 'high', 'low', 'close', 'vol']
    for col in numeric_cols:
        try:
            df[col] = pd.to_numeric(df[col], errors='coerce')
            # 填充可能的NaN值
            df[col].fillna(method='ffill', inplace=True)
            df[col].fillna(method='bfill', inplace=True)
            df[col].fillna(df[col].mean(), inplace=True)
        except Exception as e:
            print(f"转换 {col} 为数值类型时出错: {str(e)}")
            # 创建随机数据
            df[col] = np.random.uniform(10, 100, len(df))

    # 计算简单的统计特征
    # 添加简单移动平均线(SMA)
    df['ma5'] = df['close'].rolling(window=5).mean()
    df['ma10'] = df['close'].rolling(window=10).mean()
    df['ma20'] = df['close'].rolling(window=20).mean()

    # 计算涨跌幅
    df['pct_change'] = df['close'].pct_change()

    # 价格与均线差距
    df['close_ma5_diff'] = df['close'] - df['ma5']
    df['close_ma10_diff'] = df['close'] - df['ma10']

    # 波动指标
    df['vol_change'] = df['vol'].pct_change()
    df['price_range'] = (df['high'] - df['low']) / df['close']

    # 尝试添加TALib指标
    try:
        # 添加MACD
        df['macd'], df['macd_signal'], df['macd_hist'] = talib.MACD(df['close'])

        # 添加RSI
        df['rsi'] = talib.RSI(df['close'])

        # 添加布林带
        df['upper_band'], df['middle_band'], df['lower_band'] = talib.BBANDS(df['close'])

        # 添加KDJ
        df['k'], df['d'] = talib.STOCH(df['high'], df['low'], df['close'])
        df['j'] = 3 * df['k'] - 2 * df['d']

        # 添加OBV
        df['obv'] = talib.OBV(df['close'], df['vol'])
    except Exception as e:
        print(f"计算TALib指标时出错: {str(e)}")

    # 填充NaN值
    df.fillna(method='ffill', inplace=True)
    df.fillna(method='bfill', inplace=True)

    return df

def create_sequences(data, seq_length=20):
    """从数据创建序列用于预测"""
    X = []
    for i in range(len(data) - seq_length):
        X.append(data[i:i+seq_length])
    return np.array(X)

def predict_stock_trend(stock_data, model_path, feature_scaler_path=None, target_scaler_path=None, seq_length=20):
    """使用保存的模型预测趋势"""
    # 加载模型
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"模型文件不存在: {model_path}")

    model = load_model(model_path)

    # 添加技术指标
    stock_data = add_technical_indicators(stock_data)

    # 移除非数值列
    numeric_data = stock_data.select_dtypes(include=[np.number])

    # 如果有trade_date列，保存它用于结果展示
    dates = None
    if 'trade_date' in stock_data.columns:
        dates = stock_data['trade_date'].values

    # 数据归一化
    feature_scaler = MinMaxScaler()
    target_scaler = MinMaxScaler()

    # 分离特征和目标
    features = numeric_data.drop('close', axis=1, errors='ignore')
    target = numeric_data[['close']]

    # 归一化特征
    X_scaled = feature_scaler.fit_transform(features)

    # 归一化目标
    target_scaler.fit(target)

    # 创建序列
    X_seq = create_sequences(X_scaled, seq_length)

    # 如果序列为空，无法预测
    if len(X_seq) == 0:
        print("警告: 数据量不足，无法创建序列进行预测")
        return None

    # 预测
    predictions_scaled = model.predict(X_seq)

    # 反归一化预测结果
    dummy = np.zeros((len(predictions_scaled), target.shape[1]))
    dummy[:, 0] = predictions_scaled.flatten()
    predictions = target_scaler.inverse_transform(dummy)[:, 0]

    # 创建结果DataFrame
    result_dates = dates[seq_length:] if dates is not None else np.arange(len(predictions))
    actual_prices = stock_data['close'].values[seq_length:]

    results = pd.DataFrame({
        '日期': result_dates,
        '实际价格': actual_prices,
        '预测价格': predictions
    })

    # 计算预测准确率
    results['涨跌实际'] = results['实际价格'].diff() > 0
    results['涨跌预测'] = results['预测价格'].diff() > 0
    results['预测正确'] = results['涨跌实际'] == results['涨跌预测']

    accuracy = results['预测正确'].mean()
    print(f"涨跌预测准确率: {accuracy:.2%}")

    # 绘制预测结果
    plt.figure(figsize=(14, 7))
    plt.plot(results['日期'], results['实际价格'], label='实际价格', color='blue')
    plt.plot(results['日期'], results['预测价格'], label='预测价格', color='red', linestyle='--')
    plt.title('价格预测结果')
    plt.xlabel('日期')
    plt.ylabel('价格')
    plt.legend()
    plt.grid(True)

    # 如果日期太多，只显示部分刻度
    if len(results) > 20:
        plt.xticks(results['日期'][::len(results)//10])

    plt.tight_layout()
    plt.savefig('stock_prediction_result.png')
    plt.show()

    # 预测未来趋势
    last_sequence = X_scaled[-seq_length:]
    future_prediction_scaled = model.predict(np.array([last_sequence]))

    # 反归一化
    dummy = np.zeros((1, target.shape[1]))
    dummy[0, 0] = future_prediction_scaled[0, 0]
    future_prediction = target_scaler.inverse_transform(dummy)[0, 0]

    last_price = stock_data['close'].values[-1]
    trend = "上涨" if future_prediction > last_price else "下跌"
    change_pct = abs(future_prediction - last_price) / last_price * 100

    print(f"\n预测未来趋势: {trend}")
    print(f"当前价格: {last_price:.2f}")
    print(f"预测价格: {future_prediction:.2f}")
    print(f"变动百分比: {change_pct:.2f}%")

    return results

def main():
    # 模型路径
    model_path = 'saved_models/gru_lstm_model.keras'

    # 加载要预测的数据
    # 这里需要替换为实际的数据文件路径
    stock_file = input("请输入数据CSV文件路径: ")

    if not os.path.exists(stock_file):
        print(f"错误: 文件不存在 {stock_file}")
        return

    try:
        stock_data = pd.read_csv(stock_file)
        print(f"成功加载数据，共 {len(stock_data)} 行")

        # 预测趋势
        results = predict_stock_trend(stock_data, model_path)

        if results is not None:
            # 保存预测结果
            results.to_csv('prediction_results.csv', index=False)
            print(f"预测结果已保存到 prediction_results.csv")

    except Exception as e:
        print(f"预测过程中出错: {str(e)}")

if __name__ == "__main__":
    main()
